import json
import csv
import yaml
import os
from pathlib import Path


class DataLoader:

    def __init__(self):
        self.data_dir = Path(__file__).resolve().parent / 'test_data'

    def load_json(self, filename):
        """从JSON文件加载测试数据"""
        file_path = self.data_dir / filename
        with open(file_path, 'r', encoding='utf-8') as file:
            return json.load(file)

    def load_csv(self, filename):
        """从CSV文件加载测试数据"""
        file_path = self.data_dir / filename
        data = []
        with open(file_path, 'r', encoding='utf-8') as file:
            reader = csv.DictReader(file)
            for row in reader:
                data.append(row)
        return data

    def load_yaml(self, filename):
        """从YAML文件加载测试数据"""
        file_path = self.data_dir / filename
        with open(file_path, 'r', encoding='utf-8') as file:
            return yaml.safe_load(file)

    def get_test_data(self, data_type, filename, test_case_name=None):
        """
        获取测试数据
        :param data_type: 数据类型 'json', 'csv', 'yaml'
        :param filename: 文件名
        :param test_case_name: 测试用例名称（用于JSON和YAML）
        """
        if data_type == 'json':
            data = self.load_json(filename)
            return data.get(test_case_name, []) if test_case_name else data

        elif data_type == 'csv':
            return self.load_csv(filename)

        elif data_type == 'yaml':
            data = self.load_yaml(filename)
            return data.get(test_case_name, []) if test_case_name else data

        else:
            raise ValueError(f"不支持的数据类型: {data_type}")